import time
import numpy as np
import cv2
import simplejpeg
from hobot_vio import libsrcampy  # type: ignore


class HDMI_Display():
    def __init__(self, width=1024, height=600, display_id=0):
        self.width = width
        self.height = height
        self.disp = libsrcampy.Display()
        ret = self.disp.display(display_id, self.width, self.height, 0, 1)
        if ret != 0:
            raise RuntimeError(f"Display display {display_id} return: {ret}")
        print(f"Display display {display_id} initialized successfully")
    def bgr2nv12(self, image):
        """BGR to NV12 conversion"""
        height, width = image.shape[:2]
        area = height * width
        yuv420p = cv2.cvtColor(image, cv2.COLOR_BGR2YUV_I420).reshape((area * 3 // 2,))
        y = yuv420p[:area]
        uv_planar = yuv420p[area:].reshape((2, area // 4))
        uv_packed = uv_planar.transpose((1, 0)).reshape((area // 2,))
        nv12 = np.zeros_like(yuv420p)
        nv12[:area] = y
        nv12[area:] = uv_packed
        return nv12
    def display_frame(self, frame):
        """传入BGR格式的frame进行显示"""
        if frame.shape[1] != self.width or frame.shape[0] != self.height:
            frame = cv2.resize(frame, (self.width, self.height))

        nv12_img = self.bgr2nv12(frame)

        if nv12_img.size != int(self.width * self.height * 1.5):
            print("YUV尺寸不符，跳过该帧")
            return -1

        ret = self.disp.set_img(nv12_img.tobytes())
        if ret != 0:
            print(f"Display set_img return: {ret}")
        return ret

    def close(self):
        self.disp.close()
        print("Display closed")

# 使用示例
# if __name__ == "__main__":
#     cap = cv2.VideoCapture(0)
#     if not cap.isOpened():
#         print("无法打开摄像头")
#         exit(1)

#     display = HDMI_Display()

#     try:
#         while True:
#             ret, frame = cap.read()
#             if not ret:
#                 print("读取摄像头帧失败")
#                 break

#             display.display_frame(frame)
            
#     except KeyboardInterrupt:
#         print("用户中断")
#     finally:
#         cap.release()
#         display.close()
